{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GGxjzJPFoCVO"
      },
      "source": [
        "# Extending Caravan to new basins\n",
        "\n",
        "Author: Frederik Kratzert\n",
        "Contact: kratzert@google.com\n",
        "\n",
        "Last updated: Apr 13, 2024\n",
        "\n",
        "This notebook is part of the [Caravan publication](https://www.nature.com/articles/s41597-023-01975-w) and is the first of two notebooks that can be used to extend the dataset to new basins. This notebook is intended to be run on Google Colab so that the data exports to Earth Engine works as intended.\n",
        "\n",
        "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.sandbox.google.com/github/kratzert/Caravan/blob/main/code/Caravan_part1_Earth_Engine.ipynb)\n",
        "\n",
        "## What is Caravan?\n",
        "\n",
        "Caravan is ...\n",
        "- an open source, large-sample rainfall-runoff dataset.\n",
        "- derived entirely in the cloud (Earth Engine).\n",
        "- a dataset that only uses globally available data products (namely [ECMWF's ERA5-Land](https://cds.climate.copernicus.eu/cdsapp#!/dataset/reanalysis-era5-land?tab=overview) and [HydroATLAS](https://www.hydrosheds.org/page/hydroatlas) by Linke et al.).\n",
        "\n",
        "Why?\n",
        "- to democratize large-sample hydrology: Anyone can extend the dataset to new regions without downloading and preprocessing large amounts of raw data. Instead, we use freely available cloud resources to do all heavy lifting in the cloud and only download the processed results.\n",
        "- to advance research on global rainfall-runoff models: So far, we have seen  different large-sample datasets (e.g. CAMELS variants) that are specific to individual countries. However, due to the use of data products with only local coverage (e.g. meteorological forcing data or static attribute maps), intercomparisons betweens these products are not straightforward.\n",
        "\n",
        "## Requirements\n",
        "\n",
        "- You need a Google account to be able to use Earth Engine\n",
        "- You need to add your shapefile as an [Asset](https://developers.google.com/earth-engine/guides/table_upload#upload-a-shapefile) to Google Earth Engine\n",
        "- The shapefile needs to have one field that indicates the basin/gauge ID, which will be used to link the derived data to the individual basin. Make sure to adapt the corresponding variable in the \"General configuration\" section below. It is recommend to name this field `gauge_id`, which will safe you some adaptations across the 2nd notebook.\n",
        "\n",
        "Note: This notebook only derives catchment attributes and meteorological forcing data for each polygon in the shapefile layer. We do not provide additional streamflow data beyond what is already included in Caravan.\n",
        "\n",
        "## Making new data available to the community\n",
        "\n",
        "We envision Caravan as a dynamically growing dataset that can be extended by its users. While making a large-sample dataset available upon the intial publication, the distribution of the basins is still limited to a few regions in the world. Ideally, this will change over time. Anyone who uses Caravan and this code to extend the dataset is invited to also share the streamflow time series data with the community, even if it is just for a single (or a few) basins. By doing so, little by little, we could create a new, globally spanning open source dataset that anyone can use for their own work/research. For the moment, the [discussion forum](https://github.com/kratzert/caravan/discussions) on the Github repository is the place to share information on data that you want to contribute. See also the [wiki](https://github.com/kratzert/caravan/wiki) on Github for further information.\n",
        "\n",
        "##  Structure of this notebook\n",
        "\n",
        "The code consists of two different sections:\n",
        "\n",
        "1. To derive static catchment attributes from HydroATLAS.\n",
        "2. To spatially aggregate ERA5-Land forcing data for each polygon.\n",
        "\n",
        "While the first point takes a couple of seconds up to a few minutes, depending on the number of basins, the second point takes considerably longer. However, within this notebook we only start jobs to process ERA5-Land on Earth Engine and to store the data in Google Drive. The tasks will run in the background and this notebook can be closed after the end of this notebook is reached. For details, see the introduction in each of the two data processing sections.\n",
        "\n",
        "## Citation\n",
        "\n",
        "If you use the Caravan dataset or this code, please consider referencing the original publication\n",
        "\n",
        "```bib\n",
        "@article{kratzert2023caravan,\n",
        "  title={Caravan-A global community dataset for large-sample hydrology},\n",
        "  author={Kratzert, Frederik and Nearing, Grey and Addor, Nans and Erickson, Tyler and Gauch, Martin and Gilon, Oren and Gudmundsson, Lukas and Hassidim, Avinatan and Klotz, Daniel and Nevo, Sella and others},\n",
        "  journal={Scientific Data},\n",
        "  volume={10},\n",
        "  number={1},\n",
        "  pages={61},\n",
        "  year={2023},\n",
        "  publisher={Nature Publishing Group UK London}\n",
        "}\n",
        "```"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EQvLi9Bm3eRi"
      },
      "source": [
        "# General Configurations\n",
        "\n",
        "- `ASSET_NAME`: String defining the path of the asset in Google Earth Engine. Usually something like `users/{username}/{asset_name}`.\n",
        "\n",
        "- `OUTPUT_FOLDER_NAME`: The name of the folder (in Google Drive) that is used to store the resulting data. For example, `OUTPUT_FOLDER_NAME = 'my_caravan_extension'` will create a folder called `my_caravan_extension` in the main directory of your Google Drive.\n",
        "\n",
        "- `BASIN_ID_FIELD`: Name of the attribute field in the shapefile that contains the basin id. It is recommended to name this field `gauge_id`. Make sure that this name does not conflict with any of the field names of HydroATLAS. For example, you could use something like `'gauge_id'` or `'basin_id'` but _not_ `'HYBAS_ID'`, `'PFAF_ID'`, or `'MAIN_BAS'`, which are all existing field names in HydroATLAS/HydroSHEDS.\n",
        "\n",
        "- `BASIN_PREFIX`: A short descriptive string that is prepended to each basin id and that should be unique within the Caravan data space. For example, we use `camels` for basins from the CAMELS (US) dataset and `camelsgb` from the CAMELS-GB dataset. The final name for each basin within the attribute table will be `{BASIN_PREFIX}_{GAUGE_ID}`. Note, if you already included such a prefix in the basin id field of your shapefile, leave this field as an empty string. Please also read the README in the Caravan dataset folder on details about the folder structure of the dataset.\n",
        "\n",
        "- `AREA_MAX_THRESHOLD` (in km2): Defines the upper bound for the basin area. Larger basins in the shapefile are ignored when deriving the attributes and forcings. If `None`, all basins are considered. Note: To avoid user quota errors when processing too large basins, HydroATLAS attributes are derived from different HydroATLAS levels, given the size of a catchment. See [Configuration for HydroATLAS](#configuration-for-hydroatlas).\n",
        "\n",
        "- `AREA_MIN_THRESHOLD` (in km2): Defines the lower bound for the basin area. Smaller basins in the shapefile are ignored when deriving the attributes and forcings. If `None`, all basins are considered. Note: Due to the underlying resolution of HydroATLAS the attributes for basins that are much smaller than the HydroATLAS subpolygons will potentially become less and less representative of your catchment.\n",
        "\n",
        "- `EE_PROJECT_NAME`: Name of your personal Earth Engine project. Usually some string starting like `ee-{username}` if you followed the default project creation. Required to access your uploaded assets."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7FqgYvtT3e8w"
      },
      "outputs": [],
      "source": [
        "ASSET_NAME = ''\n",
        "OUTPUT_FOLDER_NAME = ''\n",
        "BASIN_ID_FIELD = 'gauge_id'\n",
        "BASIN_PREFIX = ''\n",
        "AREA_MAX_THRESHOLD = None\n",
        "AREA_MIN_THRESHOLD = None\n",
        "EE_PROJECT_NAME = ''"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PC2x5IKi4zAJ"
      },
      "source": [
        "# Import modules\n",
        "\n",
        "Note: to be able to store data in Google Drive, you need to mount your Drive in this runtime. Follow the instructions of the pop-up window, otherwise you won't be able to store the HydroATLAS attributes or ERA5-Land data.\n",
        "\n",
        "Also note: You also need to give this runtime permission to access Earth Engine, follow the instructions in the text and paste the token from the link within the box below, then press Enter. For a more detailed description you can also check our [guide](https://github.com/kratzert/Caravan/wiki/Extending-Caravan-with-new-basins)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QyQTV9jO4ylh"
      },
      "outputs": [],
      "source": [
        "from collections import defaultdict\n",
        "from pathlib import Path\n",
        "\n",
        "import ee\n",
        "import numpy as np\n",
        "import pandas as pd\n",
        "from google.colab import drive\n",
        "from tqdm.notebook import tqdm\n",
        "\n",
        "# Mount Google drive to colab runtime.\n",
        "drive.mount('/content/drive')\n",
        "drive_path = Path(f'/content/drive/My Drive/')\n",
        "\n",
        "# Authorize this Colab to access Earth Engine.\n",
        "ee.Authenticate()\n",
        "ee.Initialize(project=EE_PROJECT_NAME)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wSdmgDYo6j5K"
      },
      "source": [
        "# Utility functions\n",
        "\n",
        "A set of functions that are used below. Simply execute this block to make the functions available."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TaN32dpz6oC_"
      },
      "outputs": [],
      "source": [
        "# @title Implementation of utility function\n",
        "\n",
        "\n",
        "def add_area_as_property(poly):\n",
        "  \"\"\"Adds polygon area (km2) as property to each feature.\n",
        "\n",
        "  Different shapefiles might have different property names and some are lacking\n",
        "  the area property all together. We use this function to have a homogeneous\n",
        "  area property we can later use for filtering catchments by the area.\n",
        "  \"\"\"\n",
        "  new_poly = ee.Feature(poly).set({\"area\": poly.area().divide(1000 * 1000)})\n",
        "  return new_poly\n",
        "\n",
        "\n",
        "def subset_properties(feat, property_name):\n",
        "  \"\"\"Returns feature collection with subset of properties.\"\"\"\n",
        "  properties = feat.propertyNames()\n",
        "  selected_properties = properties.filter(ee.Filter.eq(\"item\", property_name))\n",
        "  return feat.select(selected_properties)\n",
        "\n",
        "\n",
        "def join_features(poly):\n",
        "  \"\"\"Combine two polygons and add intersecting area as property.\"\"\"\n",
        "  primary = ee.Feature(poly.get(\"primary\"))\n",
        "  secondary = ee.Feature(poly.get(\"secondary\"))\n",
        "\n",
        "  # perform intersection (creates new object)\n",
        "  newPoly = primary.intersection(secondary)\n",
        "\n",
        "  # get area of overlap\n",
        "  area = newPoly.area().divide(1000 * 1000)\n",
        "\n",
        "  # copy properties of both initial polygons in new object\n",
        "  newPoly = newPoly.copyProperties(primary).copyProperties(secondary)\n",
        "\n",
        "  # return with the area over intersection added as property\n",
        "  return ee.Feature(newPoly).set({\"Intersect\": area}).setGeometry(None)\n",
        "\n",
        "\n",
        "def start_export_task_img_coll(img_coll, feature_coll, folder, filename):\n",
        "  \"\"\"Start Export task on Earth Engine.\n",
        "\n",
        "  This function is used to send one processing and exporting task for ERA5-Land\n",
        "  data to Earth Engine.\n",
        "  \"\"\"\n",
        "  func = lambda img: img.reduceRegions(\n",
        "      collection=feature_coll, reducer=ee.Reducer.mean(), scale=5000\n",
        "  )\n",
        "  results_table = img_coll.map(func)\n",
        "  new_table = results_table.flatten().map(lambda x: x.setGeometry(None))\n",
        "\n",
        "  task = ee.batch.Export.table.toDrive(\n",
        "      collection=new_table,\n",
        "      folder=folder,\n",
        "      description=f\"{filename}_era5land\",\n",
        "      fileNamePrefix=filename,\n",
        "  )\n",
        "  task.start()\n",
        "  return task\n",
        "\n",
        "\n",
        "def compute_pour_point_properties(\n",
        "    basin_data, min_overlap_threshold, pour_point_properties\n",
        "):\n",
        "  \"\"\"Finds most downstream polygon to extract pour point features.\n",
        "\n",
        "  Some features are only defined for the entire upstream area. Here, we need to\n",
        "  find the most downstream HydroATLAS polygon that still has a significant\n",
        "  ( > 50%) overlap with the catchment polygon and take the attributes from there.\n",
        "  In case of non-significant overlap of the most downstream polygon, it is\n",
        "  possible that two (or more) merging rivers need to be considered that all drain\n",
        "  into that most downstream polygon. If that is the case, we take the sum of the\n",
        "  pour-point properties (as all of these properties are volumnes, counts, areas\n",
        "  etc.)\n",
        "  \"\"\"\n",
        "  # Find basin with the largest overlap (in percent).\n",
        "  percentage_overlap = [\n",
        "      x / y for x, y in zip(basin_data[\"weights\"], basin_data[\"SUB_AREA\"])\n",
        "  ]\n",
        "  current_basin_pos = np.argmax(percentage_overlap)\n",
        "\n",
        "  # Get the HYBAS_ID of the next downstream gauge.\n",
        "  next_down_id = basin_data[\"NEXT_DOWN\"][current_basin_pos]\n",
        "\n",
        "  # Traverse the river network downstream until we hit the termination condition.\n",
        "  while True:\n",
        "\n",
        "    # Make sure that we did not reached the ocean (HYBAS_ID=0), otherwise we would\n",
        "    # later pick all intersecting polygons that also drain into HYBAS_ID = 0.\n",
        "    if next_down_id == 0:\n",
        "      break\n",
        "\n",
        "    # Check if next_down_id is in the list of HYBAS_IDs.\n",
        "    if next_down_id not in basin_data[\"HYBAS_ID\"]:\n",
        "      break\n",
        "\n",
        "    # Get position of next_down_id in the list of HYBAS_IDs\n",
        "    next_down_pos = basin_data[\"HYBAS_ID\"].index(next_down_id)\n",
        "\n",
        "    # Check that the intersection of the next downstream polygon is at least 50%\n",
        "    if percentage_overlap[next_down_pos] < 0.5:\n",
        "      break\n",
        "\n",
        "    # Continue with the next polygon.\n",
        "    next_down_id = basin_data[\"NEXT_DOWN\"][next_down_pos]\n",
        "\n",
        "  # Find all polygons that would have drained into that polygons, which can\n",
        "  # potentially be more than one in case where two (or more) rivers merge.\n",
        "  direct_upstream_polygons = []\n",
        "  for i, next_down in enumerate(basin_data[\"NEXT_DOWN\"]):\n",
        "    if (next_down == next_down_id) and (\n",
        "        (basin_data[\"weights\"][i] > min_overlap_threshold)\n",
        "        or (basin_data[\"weights\"][i] / basin_data[\"SUB_AREA\"][i] > 0.5)\n",
        "    ):\n",
        "      direct_upstream_polygons.append(i)\n",
        "\n",
        "  # Finally compute metrics. As all pour_point_properties are volumes or\n",
        "  # areas, we can simply compute the sum.\n",
        "  aggregated_properties = {}\n",
        "  for prop in pour_point_properties:\n",
        "    aggregated_properties[prop] = sum(\n",
        "        [basin_data[prop][i] for i in direct_upstream_polygons]\n",
        "    )\n",
        "\n",
        "  return aggregated_properties\n",
        "\n",
        "\n",
        "def get_hydroatlas_intersections(\n",
        "    basin_fc: ee.FeatureCollection,\n",
        "    hydroatlas_fc: ee.FeatureCollection,\n",
        "    batch_size: int,\n",
        "    min_overlap_threshold: int,\n",
        "):\n",
        "\n",
        "  # Get the number of basins in the feature collection.\n",
        "  n_basins = basin_fc.size().getInfo()\n",
        "\n",
        "  # create Filterobject for the intersection\n",
        "  spatialFilter = ee.Filter.intersects(\n",
        "      leftField=\".geo\", rightField=\".geo\", maxError=10\n",
        "  )\n",
        "\n",
        "  # create join object\n",
        "  join = ee.Join.inner()\n",
        "\n",
        "  # list of all available basins\n",
        "  basin_ids = basin_fc.aggregate_array(BASIN_ID_FIELD).getInfo()\n",
        "\n",
        "  # dictionary to store the results\n",
        "  results = {b: defaultdict(list) for b in basin_ids}\n",
        "\n",
        "  # get number of batch requests to make\n",
        "  n_batches = n_basins // batch_size\n",
        "  if n_basins % batch_size > 0:\n",
        "    n_batches += 1\n",
        "\n",
        "  # iterate over batches\n",
        "  for batch in tqdm(range(n_batches)):\n",
        "    start_idx = batch * batch_size\n",
        "    if (batch + 1) * batch_size > n_basins:\n",
        "      process_n_basins = n_basins - start_idx\n",
        "    else:\n",
        "      process_n_basins = batch_size\n",
        "\n",
        "    # perform intersection\n",
        "    intersectJoined = join.apply(\n",
        "        ee.FeatureCollection(basin_fc.toList(process_n_basins, start_idx)),\n",
        "        hydroatlas_fc,\n",
        "        spatialFilter,\n",
        "    )\n",
        "    intersected_basins = intersectJoined.map(join_features).getInfo()\n",
        "\n",
        "    # iterate over intersected polygon and extract the information needed\n",
        "    for polygon in intersected_basins[\"features\"]:\n",
        "\n",
        "      basin_id = polygon[\"properties\"][BASIN_ID_FIELD]\n",
        "\n",
        "      # check if the intersect is larger than the minimum overlap treshold\n",
        "      if (polygon[\"properties\"][\"Intersect\"] > min_overlap_threshold) or (\n",
        "          polygon[\"properties\"][\"Intersect\"] / polygon[\"properties\"][\"SUB_AREA\"]\n",
        "          > 0.5\n",
        "      ):\n",
        "\n",
        "        # For some reason, the intersections for some basins are returned twice,\n",
        "        # leading to errors when computing the aggregates and especially the basin\n",
        "        # area. Here, we make sure that we do not add an intersect again, by\n",
        "        # filtering for unique intersect areas.\n",
        "        if (\n",
        "            polygon[\"properties\"][\"Intersect\"]\n",
        "            not in results[basin_id][\"weights\"]\n",
        "        ):\n",
        "\n",
        "          for prop in USE_PROPERTIES:\n",
        "            results[basin_id][prop].append(polygon[\"properties\"][prop])\n",
        "\n",
        "          # extract area of intersection, used as weight\n",
        "          results[basin_id][\"weights\"].append(\n",
        "              polygon[\"properties\"][\"Intersect\"]\n",
        "          )\n",
        "\n",
        "      # To compute the true basin area (that is similar to the area of the basin\n",
        "      # polygon), we need the area of all intersecting HydroATLAS polygons, even\n",
        "      # the small bits that we exclude when we aggregate the attributes.\n",
        "      if (\n",
        "          polygon[\"properties\"][\"Intersect\"]\n",
        "          not in results[basin_id][\"area_fragments\"]\n",
        "      ):\n",
        "        results[basin_id][\"area_fragments\"].append(\n",
        "            polygon[\"properties\"][\"Intersect\"]\n",
        "        )\n",
        "\n",
        "  return results\n",
        "\n",
        "\n",
        "def aggregate_hydroatlas_intersections(\n",
        "    group_results: dict[str, dict], min_overlap_threshold: float\n",
        "):\n",
        "\n",
        "  aggregated_results = defaultdict(dict)\n",
        "  for basin, basin_data in tqdm(group_results.items()):\n",
        "\n",
        "    if basin == \"camelscl_2120001\":\n",
        "      continue\n",
        "\n",
        "    # extract weights for the weighted aggregation of the attributes\n",
        "    weights = np.array(basin_data[\"weights\"])\n",
        "\n",
        "    # compute the mask of intersections with a sufficient overlap\n",
        "    mask = weights > min_overlap_threshold\n",
        "    masked_weights = weights[mask]\n",
        "\n",
        "    # perform aggregation\n",
        "    for key, val in basin_data.items():\n",
        "      # Ignore these auxilary properties\n",
        "      if key in [\n",
        "          \"weights\",\n",
        "          \"UP_AREA\",\n",
        "          \"area_fragments\",\n",
        "          \"HYBAS_ID\",\n",
        "          \"NEXT_DOWN\",\n",
        "          \"SUB_AREA\",\n",
        "      ]:\n",
        "        continue\n",
        "\n",
        "      # Pour-point properties are treated separately below.\n",
        "      if key in POUR_POINT_PROPERTIES:\n",
        "        continue\n",
        "\n",
        "      val = np.array(val)\n",
        "\n",
        "      # only consider intersections with sufficient overlap\n",
        "      masked_val = val[mask]\n",
        "\n",
        "      # for wetland classes, define 'no wetland' as a new class.\n",
        "      # In the dataset this value is -999, here we set it to 13\n",
        "      if key == \"wet_cl_smj\":\n",
        "        masked_val[masked_val == -999] = 13\n",
        "\n",
        "      # if all values are -999, set value to NaN\n",
        "      if len(masked_val[masked_val == -999]) == len(masked_val):\n",
        "        aggregated_results[basin][key] = np.nan\n",
        "      else:\n",
        "        # Majority vote for the class properties\n",
        "        if key in MAJORITY_PROPERTIES:\n",
        "          aggregated_results[basin][key] = np.bincount(\n",
        "              masked_val[masked_val > -999],\n",
        "              weights=masked_weights[masked_val > -999],\n",
        "          ).argmax()\n",
        "\n",
        "        # Averaging for all remaining properties\n",
        "        else:\n",
        "          aggregated_results[basin][key] = np.average(\n",
        "              masked_val[masked_val > -999],\n",
        "              weights=masked_weights[masked_val > -999],\n",
        "          )\n",
        "\n",
        "    aggregated_pour_point_features = compute_pour_point_properties(\n",
        "        basin_data=basin_data,\n",
        "        min_overlap_threshold=min_overlap_threshold,\n",
        "        pour_point_properties=POUR_POINT_PROPERTIES,\n",
        "    )\n",
        "    for key, val in aggregated_pour_point_features.items():\n",
        "      aggregated_results[basin][key] = val\n",
        "\n",
        "    # basin area is the sum of the area of all intersecting fragments, also\n",
        "    # the smaller intersections that are ignored above.\n",
        "    aggregated_results[basin][\"area\"] = sum(basin_data[\"area_fragments\"])\n",
        "\n",
        "    # We also store the fraction of the area considered during aggregation\n",
        "    # to the total basin area.\n",
        "    aggregated_results[basin][\"area_fraction_used_for_aggregation\"] = sum(\n",
        "        masked_weights\n",
        "    ) / sum(basin_data[\"area_fragments\"])\n",
        "\n",
        "  return aggregated_results"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WdNCyuqa1wiz"
      },
      "source": [
        "# Processing HydroATLAS\n",
        "\n",
        "This section of the code derives the selected attributes from HydroATLAS that are considered in Caravan for every polygon in the shape layer.\n",
        "\n",
        "## TLDR; How are the attributes derived?\n",
        "\n",
        "- We use HydroATLAS level 12, i.e. the layer with the highest available resolution.\n",
        "- Next, we take the polygons of the catchment boundaries and for each polygon, we compute the intersecting polygon areas of the HydroATLAS dataset.\n",
        "- With the resulting clipped polygons, we derive the attributes as aggregates (see the \"HydroATLAS features\" sub-section for more details on the aggregation procedure).\n",
        "\n",
        "## Output\n",
        "\n",
        "After completing this part of the code, you will find an `attributes_new.csv` file in your predefined Google Drive folder."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ABIHP_vr59Lt"
      },
      "source": [
        "## Configurations for HydroATLAS\n",
        "\n",
        "HydroATLAS polygons exist in different resolution, with level 12 being the highest resolution and the smaller the level the coarser the resolution. When deriving attributes, we intersect the catchment polygons with the HydroATLAS polygons, leading to potentially a very large number of intersecting HydroATLAS attributes for larger basins. This can result in user quota errors on Earth Engine and arguably it is also not making a real difference is larger catchments have their attributes derived from the finest resolution of subpolygons or already use any of the coarser resolutions (that are pre-aggregated from level 12). When processing the Caravan dataset, the configuration below was used.\n",
        "\n",
        "Note: We manually compared basins attributes derived from the neighboring configuration (e.g. catchments larger than 50k derived from HydroATLAS level 12) and could not find that it makes a real difference. However, you are free to adapt the code below. Note that if you want to use the finest resolution possible, you might need to decrease the batch size.\n",
        "\n",
        "- `min_overlap_threshold`:  (in km2): Because of differences in the basin delinations between HydroSHEDS and the user-provided catchment boundaries, the outer border of a catchment does not necessarily perfectly match the boundaries of a HydroSHEDS polygon. This can lead to some artifacts during the intersection process, where small pieces of a neighboring HydroSHEDS polygon are included. We use the value of this variable to filter these small artifacts and only consider overlapping polygons that have an overlapping area that is greater than `min_overlap_threshold`. However, this means that this variable also marks a minimum catchment size that can be processed with this code.\n",
        "\n",
        "- `batch_size`: This value defines the number of basins that are processed in one call to the Earth Engine API. Generally, a higher value is better, since Earth Engine is extremely good in parallelizing jobs, but there is an upper limit of how many objects can be processed at once. A value of 50 was a reasonable tradeoff and seemed to not yield any problems. In case you get an error below when processing your basins (something like \"EEException: Collection query aborted after accumulating over 5000 elements.\"), try reducing the batch size.\n",
        "\n",
        "- `hydroatlas`: Asset ID of the HydroATLAS layer to use for this configuration. E.g. `\"WWF/HydroATLAS/v1/Basins/level12\"` for level 12.\n",
        "\n",
        "- `area_greater_than`: Lower area threshold (exclusive). `None` for no lower threshold.\n",
        "\n",
        "- `area_not_greater_than`: Upper area threshold (inclusive). `None` for no upper threshold.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Yt6zHStNJ2dP"
      },
      "outputs": [],
      "source": [
        "CONFIGURATION = {\n",
        "    '0-to-2000': {\n",
        "        'min_overlap_threshold': 0,\n",
        "        'batch_size': 5,\n",
        "        'hydroatlas': 'WWF/HydroATLAS/v1/Basins/level12',\n",
        "        'area_greater_than': None,\n",
        "        'area_not_greater_than': 2000,\n",
        "    },\n",
        "    '2000-to-50000': {\n",
        "        'min_overlap_threshold': 5,\n",
        "        'batch_size': 10,\n",
        "        'hydroatlas': 'WWF/HydroATLAS/v1/Basins/level12',\n",
        "        'area_greater_than': 2000,\n",
        "        'area_not_greater_than': 50000,\n",
        "    },\n",
        "    '50000-to-100000': {\n",
        "        'min_overlap_threshold': 5,\n",
        "        'batch_size': 10,\n",
        "        'hydroatlas': 'WWF/HydroATLAS/v1/Basins/level09',\n",
        "        'area_greater_than': 50000,\n",
        "        'area_not_greater_than': 100000,\n",
        "    },\n",
        "    '100000-to-1000000': {\n",
        "        'min_overlap_threshold': 5,\n",
        "        'batch_size': 5,\n",
        "        'hydroatlas': 'WWF/HydroATLAS/v1/Basins/level07',\n",
        "        'area_greater_than': 100000,\n",
        "        'area_not_greater_than': 1000000,\n",
        "    },\n",
        "    '1000000-to-inf': {\n",
        "        'min_overlap_threshold': 5,\n",
        "        'batch_size': 1,\n",
        "        'hydroatlas': 'WWF/HydroATLAS/v1/Basins/level07',\n",
        "        'area_greater_than': 1000000,\n",
        "        'area_not_greater_than': None,\n",
        "    },\n",
        "}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WN6ioEX16RUv"
      },
      "source": [
        "## Initialize Feature collections\n",
        "\n",
        "The next section initializes the basin shapefile uploaded to Earth Engine as well as HydroATLAS level 12. If area thresholds are provided, this is the place where basin polygons are removed that are outside of the thresholds. The HydroATLAS layer is used to get the names of all attributes."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "H1wxbHIK6HB1"
      },
      "outputs": [],
      "source": [
        "# initialize featurecollection from uploaded asset\n",
        "basins = ee.FeatureCollection(ASSET_NAME)\n",
        "\n",
        "# Add area as polygon property and filter using the threshold defined above\n",
        "basins = basins.map(add_area_as_property)\n",
        "if AREA_MIN_THRESHOLD is not None:\n",
        "  basins = basins.filterMetadata('area', 'greater_than', AREA_MIN_THRESHOLD)\n",
        "if AREA_MAX_THRESHOLD is not None:\n",
        "  basins = basins.filterMetadata('area', 'not_greater_than', AREA_MAX_THRESHOLD)\n",
        "\n",
        "# Instantiate HydroATLAS feature collection\n",
        "hydro_atlas = ee.FeatureCollection('WWF/HydroATLAS/v1/Basins/level12')\n",
        "\n",
        "# Print number of basins that will be processed\n",
        "n_basins = basins.size().getInfo()\n",
        "print(\n",
        "    'Number of basins in feature collection and within an area bounds:'\n",
        "    f' {n_basins}'\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3suZ9_O361nS"
      },
      "source": [
        "## Filter list of HydroATLAS attributes\n",
        "\n",
        "The following cell defines which features are considered and which are ignored. Further, we define how the HydroATLAS attributes are aggregated, as there are often multiple HydroATLAS polygons within one catchment polygon due to the relatively high resolution of the level 12 layer. For discrete classes, we take the area-weighted majority vote, i.e., the class that is represented the most across all intersecting HydroATLAS polygons weighted by the area. For most of the other attributes which are defined per sub-polygon in HydroATLAS, we consider the polygon values and compute the area-weighted average. Some attributes are only defined for the entire upstream catchment. For these, we take the `UP_AREA` (upstream area) property to identify the most downstream HydroATLAS polygon and then take the attribute from here.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "fdwv6IMh6sdf",
        "outputId": "5b3c4f17-5d6b-4f97-d1d0-2ebc6ab7929f"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Found 295 features initially\n",
            "Number of ignored properties: 97 \n",
            "--Ignored 11 HydroRIVER properties\n",
            "--Ignored 86 upstream basin properties\n",
            "Remaining number of features: 194\n"
          ]
        }
      ],
      "source": [
        "# Use any HydroATLAS level to get the feature names.\n",
        "hydro_atlas = ee.FeatureCollection('WWF/HydroATLAS/v1/Basins/level12')\n",
        "\n",
        "# get list list of all available hydroatlas feature properties\n",
        "property_names = hydro_atlas.first().propertyNames().getInfo()\n",
        "\n",
        "print(f'Found {len(property_names)} features initially')\n",
        "\n",
        "# List of features for which we take the area-weighted majority vote\n",
        "MAJORITY_PROPERTIES = [\n",
        "    'clz_cl_smj',  # climate zones (18 classes)\n",
        "    'cls_cl_smj',  # climate strata (125 classes)\n",
        "    'glc_cl_smj',  # land cover (22 classes)\n",
        "    'pnv_cl_smj',  # potential natural vegetation (15 classes)\n",
        "    'wet_cl_smj',  # wetland (12 classes)\n",
        "    'tbi_cl_smj',  # terrestrial biomes (14 classes)\n",
        "    'tec_cl_smj',  # Terrestrial Ecoregions (846 classes)\n",
        "    'fmh_cl_smj',  # Freshwater Major Habitat Types (13 classes)\n",
        "    'fec_cl_smj',  # Freshwater Ecoregions (426 classes)\n",
        "    'lit_cl_smj',  # Lithological classes (16 classes)\n",
        "]\n",
        "\n",
        "# List of features for which we take the value of the most downstream polygon\n",
        "POUR_POINT_PROPERTIES = [\n",
        "    'dis_m3_pmn',  # natural discharge annual mean\n",
        "    'dis_m3_pmx',  # natural discharge annual max\n",
        "    'dis_m3_pyr',  # natural discharge annual min\n",
        "    'lkv_mc_usu',  # Lake Volume\n",
        "    'rev_mc_usu',  # reservoir volume\n",
        "    'ria_ha_usu',  # River area\n",
        "    'riv_tc_usu',  # River volumne\n",
        "    'pop_ct_usu',  # Population count in upstream area\n",
        "    'dor_pc_pva',  # Degree of regulation in upstream area\n",
        "]\n",
        "\n",
        "# These HydroSHEDS/RIVERS features are ignored\n",
        "ignore_properties = [\n",
        "    'system:index',\n",
        "    'COAST',\n",
        "    'DIST_MAIN',\n",
        "    'DIST_SINK',\n",
        "    'ENDO',\n",
        "    'MAIN_BAS',\n",
        "    'NEXT_SINK',\n",
        "    'ORDER_',\n",
        "    'PFAF_ID',\n",
        "    'SORT',\n",
        "]\n",
        "\n",
        "# These features are required to find the most downstream gauge and will be removed later.\n",
        "additional_properties = ['HYBAS_ID', 'NEXT_DOWN', 'SUB_AREA', 'UP_AREA']\n",
        "\n",
        "# These HydroATLAS features are ignored, mostly because they have a per-polygon\n",
        "# counterpart that we include.\n",
        "upstream_properties = [\n",
        "    'aet_mm_uyr',  # Actual evapotranspiration\n",
        "    'ari_ix_uav',  # Global aridity index\n",
        "    'cly_pc_uav',  # clay fraction soil\n",
        "    'cmi_ix_uyr',  # Climate Moisture Index\n",
        "    'crp_pc_use',  # Cropland Extent\n",
        "    'ele_mt_uav',  # Elevtion\n",
        "    'ero_kh_uav',  # Soil erosion\n",
        "    'for_pc_use',  # Forest cover extent\n",
        "    'gdp_ud_usu',  # Gross Domestic Product\n",
        "    'gla_pc_use',  # Glacier Extent\n",
        "    'glc_pc_u01',  # Land cover extent percent per class (22)\n",
        "    'glc_pc_u02',\n",
        "    'glc_pc_u03',\n",
        "    'glc_pc_u04',\n",
        "    'glc_pc_u05',\n",
        "    'glc_pc_u06',\n",
        "    'glc_pc_u07',\n",
        "    'glc_pc_u08',\n",
        "    'glc_pc_u09',\n",
        "    'glc_pc_u10',\n",
        "    'glc_pc_u11',\n",
        "    'glc_pc_u12',\n",
        "    'glc_pc_u13',\n",
        "    'glc_pc_u14',\n",
        "    'glc_pc_u15',\n",
        "    'glc_pc_u16',\n",
        "    'glc_pc_u17',\n",
        "    'glc_pc_u18',\n",
        "    'glc_pc_u19',\n",
        "    'glc_pc_u20',\n",
        "    'glc_pc_u21',\n",
        "    'glc_pc_u22',\n",
        "    'hft_ix_u09',  # Human Footprint 2009\n",
        "    'hft_ix_u93',  # Human Footprint 1993\n",
        "    'inu_pc_ult',  # inundation extent long-term maximum\n",
        "    'inu_pc_umn',  # inundation extent annual minimum\n",
        "    'inu_pc_umx',  # inundation extent annual maximum\n",
        "    'ire_pc_use',  # Irrigated Area Extent (Equipped)\n",
        "    'kar_pc_use',  # Karst Area Extent\n",
        "    'lka_pc_use',  # Limnicity (Percent Lake Area)\n",
        "    'nli_ix_uav',  # Nighttime Lights\n",
        "    'pac_pc_use',  # Protected Area Extent\n",
        "    'pet_mm_uyr',  # Potential evapotranspiration\n",
        "    'pnv_pc_u01',  # potential natural vegetation (15 classes)\n",
        "    'pnv_pc_u02',\n",
        "    'pnv_pc_u03',\n",
        "    'pnv_pc_u04',\n",
        "    'pnv_pc_u05',\n",
        "    'pnv_pc_u06',\n",
        "    'pnv_pc_u07',\n",
        "    'pnv_pc_u08',\n",
        "    'pnv_pc_u09',\n",
        "    'pnv_pc_u10',\n",
        "    'pnv_pc_u11',\n",
        "    'pnv_pc_u12',\n",
        "    'pnv_pc_u13',\n",
        "    'pnv_pc_u14',\n",
        "    'pnv_pc_u15',\n",
        "    'pop_ct_ssu',  # population count\n",
        "    'ppd_pk_uav',  # population density\n",
        "    'pre_mm_uyr',  # precipitation\n",
        "    'prm_pc_use',  # Permafrost extent\n",
        "    'pst_pc_use',  # Pasture extent\n",
        "    'ria_ha_ssu',  # river area in sub polygon\n",
        "    'riv_tc_ssu',  # River volumne in sub polygon\n",
        "    'rdd_mk_uav',  # Road density\n",
        "    'slp_dg_uav',  # slope degree\n",
        "    'slt_pc_uav',  # silt fraction\n",
        "    'snd_pc_uav',  # sand fraction\n",
        "    'snw_pc_uyr',  # snow cover percent\n",
        "    'soc_th_uav',  # organic carbon content in soil\n",
        "    'swc_pc_uyr',  # soil water content\n",
        "    'tmp_dc_uyr',  # air temperature\n",
        "    'urb_pc_use',  # urban extent\n",
        "    'wet_pc_u01',  # wetland classes percent (9 classes)\n",
        "    'wet_pc_u02',\n",
        "    'wet_pc_u03',\n",
        "    'wet_pc_u04',\n",
        "    'wet_pc_u05',\n",
        "    'wet_pc_u06',\n",
        "    'wet_pc_u07',\n",
        "    'wet_pc_u08',\n",
        "    'wet_pc_u09',\n",
        "    'wet_pc_ug1',  #  wetland classes percent by class grouping (2 classes)\n",
        "    'wet_pc_ug2',\n",
        "    'gad_id_smj',  # global administrative areas (country id's)\n",
        "]\n",
        "\n",
        "# get the final list of features to use (UP_AREA is only used for processing and deleted later\n",
        "# so we subtract one.)\n",
        "USE_PROPERTIES = [\n",
        "    p\n",
        "    for p in property_names\n",
        "    if p not in ignore_properties + upstream_properties\n",
        "]\n",
        "\n",
        "print(\n",
        "    'Number of ignored properties:'\n",
        "    f' {len(ignore_properties) + len(upstream_properties) + 1} '\n",
        ")\n",
        "print(f'--Ignored {len(ignore_properties) + 1} HydroRIVER properties')\n",
        "print(f'--Ignored {len(upstream_properties)} upstream basin properties')\n",
        "print(\n",
        "    'Remaining number of features:'\n",
        "    f' {len(USE_PROPERTIES) - 1 - len(additional_properties)}'\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6-7FVKUt8J3G"
      },
      "source": [
        "## Perform Intersection and download data\n",
        "\n",
        "The following piece of code queries Earth Engine to perform the intersection between the user-provided polygons and HydroATLAS. The results are unpacked into a nested dictionary. Each basin within this dictionary is another dictionary with one key per attribute and the corresponding value is a list of attribute values of all intersecting HydroATLAS polygons. Additionally, we store the area of the intersection between an individual HydroATLAS polygon and the user-provided polygons in a list to use them as aggregation weights in the next step."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "O788NrIrwYBZ"
      },
      "outputs": [],
      "source": [
        "results = {}\n",
        "\n",
        "for group_name, configuration in CONFIGURATION.items():\n",
        "  basins = ee.FeatureCollection(ASSET_NAME)\n",
        "  if configuration['area_greater_than'] is not None:\n",
        "    basins = basins.filterMetadata(\n",
        "        'area', 'greater_than', configuration['area_greater_than']\n",
        "    )\n",
        "  if configuration['area_not_greater_than'] is not None:\n",
        "    basins = basins.filterMetadata(\n",
        "        'area', 'not_greater_than', configuration['area_not_greater_than']\n",
        "    )\n",
        "  n_basins = basins.size().getInfo()\n",
        "\n",
        "  print(f'There are {n_basins} in the {group_name} group.')\n",
        "\n",
        "  if n_basins > 0:\n",
        "\n",
        "    hydroatlas_fc = ee.FeatureCollection(configuration['hydroatlas'])\n",
        "\n",
        "    group_results = get_hydroatlas_intersections(\n",
        "        basin_fc=basins,\n",
        "        hydroatlas_fc=hydroatlas_fc,\n",
        "        batch_size=configuration['batch_size'],\n",
        "        min_overlap_threshold=configuration['min_overlap_threshold'],\n",
        "    )\n",
        "\n",
        "    aggregated_results = aggregate_hydroatlas_intersections(\n",
        "        group_results=group_results,\n",
        "        min_overlap_threshold=configuration['min_overlap_threshold'],\n",
        "    )\n",
        "\n",
        "    results.update(aggregated_results)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cMVR4PtS85mI"
      },
      "source": [
        "## Create DataFrame and show results\n",
        "\n",
        "Here, we convert the dictionary into a DataFrame and check for NaN values within the data."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SL3lRYru817T"
      },
      "outputs": [],
      "source": [
        "# store all data as pandas DataFrame\n",
        "if BASIN_PREFIX:\n",
        "    df = pd.DataFrame([val for val in results.values()],\n",
        "                      index=[f\"{BASIN_PREFIX}_{str(k)}\" for k in results.keys()])\n",
        "else:\n",
        "    df = pd.DataFrame([val for val in results.values()],\n",
        "                       index=[str(k) for k in results.keys()])\n",
        "\n",
        "df.index.name = \"gauge_id\"\n",
        "df = df.sort_index(axis=0)\n",
        "\n",
        "# check for features containing NaN's\n",
        "nan_columns = df.columns[df.isna().any()].tolist()\n",
        "\n",
        "if nan_columns:\n",
        "    print(f\"{len(nan_columns)} Columns contain NaNs\")\n",
        "    print(f\"These columns contain NaN values {nan_columns}\")\n",
        "else:\n",
        "    print(\"No column contains any NaN value\")\n",
        "\n",
        "df"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8hmtF5hZ9HC2"
      },
      "source": [
        "## Save results to Drive\n",
        "\n",
        "We create an `attributes` folder within the user-defined output path into which the resulting csv file is stored."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "7wTSvkFn9CI4",
        "outputId": "65587ffb-eb53-463b-890c-5c52fd7ae804"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "HydroATLAS attributes processed successfully\n"
          ]
        }
      ],
      "source": [
        "# create output folder\n",
        "attribute_dir = drive_path / OUTPUT_FOLDER_NAME / 'attributes'\n",
        "if not attribute_dir.is_dir():\n",
        "    attribute_dir.mkdir(parents=True)\n",
        "\n",
        "# save dataframe as csv\n",
        "df.to_csv(attribute_dir / 'attributes.csv')\n",
        "\n",
        "print(\"HydroATLAS attributes processed successfully\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZBtxX7dj-p2V"
      },
      "source": [
        "# Processing ERA5-Land\n",
        "\n",
        "This section of the code prepares and starts tasks to process hourly ERA5-Land data for each basin in the shapefile. After this section is completed, you can close this notebook since the tasks were sent to Earth Engine and Earth Engine will store the resulting files in Drive.\n",
        "\n",
        "**Note**: It might take several hours until you find the processed data in your Google Drive. The time depends on the length of ERA5-Land data that you request, the number of basins in your shape layer, the size of the basins and on the numbers of points that span your polygons. If your basin polygons have a very fine resolution, consider simplifying them to speed up this step, e.g., using the [QGIS simplify function](https://docs.qgis.org/2.8/en/docs/user_manual/processing_algs/qgis/vector_geometry_tools/simplifygeometries.html) (we used the default tolerance of 1.0 for e.g. CAMELS US and CAMELS CL)\n",
        "\n",
        "## TLDR; What do we do here?\n",
        "\n",
        "Since most streamflow records are stored in local time but ERA5-Land is provided in UTC, we use ERA5-Land in hourly resolution, even if we are only interested in daily aggregates. By using hourly data, we can shift the forcing data to local time before we compute the daily aggregates and thus make sure that the forcing data aligns with the underlying streamflow data. However, this is part of the second processing step of the meteorological forcing data that runs on your local machine (no worries, this is something that any laptop/PC can do), after we used Earth Engine to compute the spatial averages.\n",
        "\n",
        "## Output\n",
        "\n",
        "The output of this part of the code is a number (we split the time into manageable slices) of csv files. Each csv file contains one row of data per basin and timestep. The data are the spatially averaged ERA5-Land variables that we request Earth Engine to process. Note: These files can become large since this is still a lot of data. If you do not have enough free space in your Google Drive, either upgrade your disk space or process the ERA5-Land time period in chunks and download the processed data after every step.\n",
        "\n",
        "After you have downloaded all files to your local machine, use our other [notebook]() to process the data locally into time series data and into the Caravan native format.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "v-S7hotuC96v"
      },
      "source": [
        "## Configuration for ERA5-Land\n",
        "\n",
        "- `IMGS_PER_REQUEST`: This defines the number of ERA5-Land images (one image = one hour) that we want to process in one request. Earth Engine is good in parallelizing compute, so a larger number is preferable. However, Earth Engine also has a size limit for a single user request. If you see an error when sending the requests that states that the query is aborted due to too many elements, consider reducing this number.\n",
        "\n",
        "- `START_DATE`: The first day of the period for which you request data. In Caravan, we include forcing data starting January 1st, 1981.\n",
        "\n",
        "- `END_DATE`: The last day of the period for which you request data. In Caravan, we include forcing data until December 31st, 2020. Note, that this date refers to data in UTC. To be able to include December 31st, 2020 in the local time zone you might actually need to request one additional day, depending on the location.\n",
        "\n",
        " `ERA5L_BANDS`: The list of bands for which you wish to compute the spatial averages. You can change this at will, however to be compatible with Caravan, you should probably leave it as is. The names have to match the names in [this table](https://developers.google.com/earth-engine/datasets/catalog/ECMWF_ERA5_LAND_HOURLY#bands). **Note** The second notebook that processes the output of Earth Engine assumes that these bands are available. If you make changes here, you might need to adapt parts of the second notebook.\n",
        "\n",
        " **Note**: Later, in the second notebook, we shift the ERA5-Land data from UTC-0 to local time, before aggregating to daily resolution. Here, we only consider days with values for all 24 hours. Therefore, adapt the `START_DATE` and the `END_DATE` according to the local timezone of your data with a 1 day buffer. That is, if your basins are east of UTC-0, adapt your `START_DATE` to include e.g. the 31st December. If the basins are west of UTC-0, adapt the `END_DATE` to include 1st January. If you want to be on the safe side, simply include a 1 day buffer on both side and then crop the data in the 2nd notebook."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fVjqXIne-Tqk"
      },
      "outputs": [],
      "source": [
        "IMGS_PER_REQUEST = 365 * 24\n",
        "START_DATE = '1950-01-01'  # Format: YYYY-MM-DD\n",
        "END_DATE = '2023-01-01'  # Format: YYYY-MM-DD\n",
        "\n",
        "# ------------------------------ Change with caution --------------------------#\n",
        "ERA5L_BANDS = [\n",
        "    'dewpoint_temperature_2m',\n",
        "    'temperature_2m',\n",
        "    'volumetric_soil_water_layer_1',\n",
        "    'volumetric_soil_water_layer_2',\n",
        "    'volumetric_soil_water_layer_3',\n",
        "    'volumetric_soil_water_layer_4',\n",
        "    'surface_net_solar_radiation',\n",
        "    'surface_net_thermal_radiation',\n",
        "    'u_component_of_wind_10m',\n",
        "    'v_component_of_wind_10m',\n",
        "    'surface_pressure',\n",
        "    'total_precipitation',\n",
        "    'snow_depth_water_equivalent',\n",
        "    'potential_evaporation',\n",
        "]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QinE2uM2EJSv"
      },
      "source": [
        "## Initialize feature collections"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "WeTlP-8JEFmR",
        "outputId": "b7a53384-efe8-4d88-ef97-c1090795be9a"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Number of basins: 1\n",
            "Number of hourly ERA5 images: 639911\n"
          ]
        }
      ],
      "source": [
        "# initialize featurecollection from uploaded asset\n",
        "basins = ee.FeatureCollection(ASSET_NAME)\n",
        "\n",
        "# Add area as polygon property and filter using the threshold defined above\n",
        "basins = basins.map(add_area_as_property)\n",
        "if AREA_MIN_THRESHOLD is not None:\n",
        "  basins = basins.filterMetadata('area', 'greater_than', AREA_MIN_THRESHOLD)\n",
        "if AREA_MAX_THRESHOLD is not None:\n",
        "  basins = basins.filterMetadata('area', 'not_greater_than', AREA_MAX_THRESHOLD)\n",
        "\n",
        "# remove all unnecessary fields from the basin polygons to reduce the export size\n",
        "basins = basins.map(lambda f: subset_properties(f, BASIN_ID_FIELD))\n",
        "\n",
        "# initialize featurecollection of ERA5-Land hourly data\n",
        "era5_collection = ee.ImageCollection('ECMWF/ERA5_LAND/HOURLY').filterDate(\n",
        "    START_DATE, END_DATE\n",
        ")\n",
        "\n",
        "# subset ERA5 bands\n",
        "era5_collection = era5_collection.select(ERA5L_BANDS)\n",
        "\n",
        "# Print some meta information\n",
        "n_basins = basins.size().getInfo()\n",
        "n_images = era5_collection.size().getInfo()\n",
        "\n",
        "print(f'Number of basins: {n_basins}')\n",
        "print(f'Number of hourly ERA5 images: {n_images}')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "q7yMLA6_F0XO"
      },
      "source": [
        "## Start processing/exporting tasks\n",
        "\n",
        "This block will create the Earth Engine jobs to process the ERA5-Land data. After this block is completed, you can close this notebook and wait for Earth Engine to store the data in your drive (might take time!). You can see the list of current tasks in the [code editor of Earth Engine](https://code.earthengine.google.com/) (right side under \"Tasks\") or use the code block below to print the current list of jobs (and their state) in this notebook."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5AnW5xDgESx8"
      },
      "outputs": [],
      "source": [
        "start = 0\n",
        "batch_number = 1\n",
        "# how many images to process in one task\n",
        "\n",
        "n_imgs = IMGS_PER_REQUEST\n",
        "while start < n_images:\n",
        "  # clip the number of images to process by the number of total available images\n",
        "  if (start + n_imgs) > n_images:\n",
        "    n_imgs = n_images - start\n",
        "\n",
        "  print(f\"Processing {n_imgs} images, starting from index {start}\")\n",
        "\n",
        "  # start export task for 'n_imgs' elements, starting at index 'start'\n",
        "  test_task = start_export_task_img_coll(\n",
        "      img_coll=ee.ImageCollection(\n",
        "          era5_collection.toList(count=n_imgs, offset=start)\n",
        "      ),\n",
        "      feature_coll=basins,\n",
        "      folder=OUTPUT_FOLDER_NAME,\n",
        "      filename=f\"batch{str(batch_number).zfill(2)}\",\n",
        "  )\n",
        "  start += n_imgs\n",
        "  batch_number += 1\n",
        "\n",
        "print(\"All jobs were sent to Earth Engine.\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HFv3cfXqImAk"
      },
      "source": [
        "# Check current tasks on Earth Engine\n",
        "\n",
        "The next section can be used optionally to check the state of the current tasks and can also be run separately of the other code (e.g., if you want to check the state of your ERA5-Land exporting tasks after a couple of hours)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "r_z17GryFqT2"
      },
      "outputs": [],
      "source": [
        "task_prefix = 'batch'\n",
        "\n",
        "\n",
        "def get_task_status(_prefix=None):\n",
        "  task_status_list = [t.status() for t in ee.batch.Task.list()]\n",
        "\n",
        "  # If a prefix was specified, use it to filter the task list.\n",
        "  if _prefix:\n",
        "    task_status_list = [\n",
        "        s for s in task_status_list if s['description'].startswith(_prefix)\n",
        "    ]\n",
        "  return task_status_list\n",
        "\n",
        "\n",
        "df = pd.DataFrame(get_task_status(task_prefix))\n",
        "\n",
        "# Convert datetime fields.\n",
        "df['creation_timestamp_ms'] = pd.to_datetime(\n",
        "    df['creation_timestamp_ms'], unit='ms', utc=True\n",
        ")\n",
        "df['start_timestamp_ms'] = pd.to_datetime(\n",
        "    df['start_timestamp_ms'], unit='ms', utc=True\n",
        ")\n",
        "df['update_timestamp_ms'] = pd.to_datetime(\n",
        "    df['update_timestamp_ms'], unit='ms', utc=True\n",
        ")\n",
        "\n",
        "df"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "last_runtime": {
        "build_target": "//intelligence/flood_forecasting/colab:notebook",
        "kind": "private"
      },
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3.10.7 64-bit",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "name": "python",
      "version": "3.10.9"
    },
    "vscode": {
      "interpreter": {
        "hash": "916dbcbb3f70747c44a77c7bcd40155683ae19c65e1c03b4aa3499c5328201f1"
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
